import numpy as np
from scipy import signal
# @save
def vec_signal(s, SNR, theta_d, k):
    '''
    单矢量水听器输出信号\\
    `s`信号源波形\\
    `SNR`信噪比\\
    `theta_d`信源方位

    '''

    L = s.shape[1]
    n=s.shape[0]
    As = np.sqrt(10**(SNR/10))
    S=As*s
    X = np.zeros((6, L),dtype=complex)
    u=np.array([np.ones((1,n)),np.cos(theta_d),np.sin(theta_d),np.cos(theta_d)**2,np.sin(theta_d)**2,np.cos(theta_d)*np.sin(theta_d)]).reshape((6,n))
    X=u@S
    noise_p = (np.random.randn(1, L)+np.random.randn(1, L)*1j)/np.sqrt(2)  # 声压通道噪声
    noise_v = (np.random.randn(3, L)+np.random.randn(3, L)*1j)/np.sqrt(2)  # 振速通道噪声
    noise_g = (np.random.randn(6,L)+np.random.randn(6,L)*1j)/np.sqrt(2) #振速梯度通道噪声
    X[0, :] = X[0,:]+noise_p  # 声压通道输出
    X[1,:]=(X[1,:]+noise_v[0,:])
    X[2,:]=(X[2,:]+noise_v[1,:])
    X[3,:]=(X[3,:]+noise_g[0,:])
    X[4,:]=(X[4,:]+noise_g[1,:])
    X[5,:]=(X[5,:]+noise_g[2,:])
    
    return X[0:k,:]

def st2_value(X):
    '''
    计算二阶统计量
    '''
    L=X.shape[1]
    pp=np.abs(sum(np.multiply(X[0,:],X[0,:])))/L
    xx=np.abs(sum(np.multiply(X[1,:],X[1,:])))/L
    yy=np.abs(sum(np.multiply(X[2,:],X[2,:])))/L
    px=np.abs(sum(np.multiply(X[0,:],X[1,:])))/L
    py=np.abs(sum(np.multiply(X[0,:],X[2,:])))/L
    xy=np.abs(sum(np.multiply(X[1,:],X[2,:])))/L
    
    return pp,xx,yy,px,py,xy
def st2_value2(X):
    '''
    计算二阶矢量水听器的二阶统计量
    '''
    L=X.shape[1]
    pgxx=np.abs(sum(np.multiply(X[0,:],X[3,:]))/L)
    pgyy=np.abs(sum(np.multiply(X[0,:],X[4,:]))/L)
    pgxy=np.abs(sum(np.multiply(X[0,:],X[5,:]))/L)
    vxgxx=np.abs(sum(np.multiply(X[1,:],X[3,:]))/L)
    vxgyy=np.abs(sum(np.multiply(X[1,:],X[4,:]))/L)
    vxgxy=np.abs(sum(np.multiply(X[1,:],X[5,:]))/L)
    vygxx=np.abs(sum(np.multiply(X[2,:],X[3,:]))/L)
    vygyy=np.abs(sum(np.multiply(X[2,:],X[4,:]))/L)
    vygxy=np.abs(sum(np.multiply(X[2,:],X[5,:]))/L)
    gxxgyy=np.abs(sum(np.multiply(X[3,:],X[4,:]))/L)
    gxxgxy=np.abs(sum(np.multiply(X[3,:],X[5,:]))/L)
    gyygxy=np.abs(sum(np.multiply(X[4,:],X[5,:]))/L)

    return pgxx, pgyy, pgxy, vxgxx, vxgyy, vxgxy, vygxx, vygyy, vygxy, gxxgyy, gxxgxy, gyygxy


def sointen_t(X):
    '''
    时域声强器
    '''
    L=X.shape[1]
    Ix=np.abs(sum(np.multiply(X[0,:],X[1,:])))/L
    Iy=np.abs(sum(np.multiply(X[0,:],X[2,:])))/L
    theta=np.arctan2(Iy,Ix)

    return theta

def sointen_f(X,fs):
    '''
    互谱声强法\\
    `X`单矢量水听器输出\\
    `fs`采样频率
    '''
    f,cpx=signal.csd(X[0,:],X[1,:],fs=fs,nperseg=1024)
    f,cpy=signal.csd(X[0,:],X[2,:],fs=fs,nperseg=1024)
    Ix=abs(sum(cpx))
    Iy=abs(sum(cpy))
    theta=np.arctan2(Iy,Ix)

    return theta
